"""
工作流管理器 - LangGraph主图架构

统一管理所有工作流，使用LangGraph StateGraph实现：
- 意图识别和路由
- 业务工作流子图集成（债券分析、对比、检索）
- 对话处理
- 记忆管理
- 完整的事件流传播（主图+子图）
"""

from typing import TypedDict, Annotated, Dict, Any, Optional, AsyncIterator, List
from datetime import datetime
import operator
import re

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.config import get_stream_writer
from langchain_core.runnables import RunnableConfig

from agents.general_intent_recognition_agent import GeneralIntentRecognitionAgent
from agents.general_conversational_agent import GeneralConversationalAgent

from workflows.bond_analysis_workflow import BondAnalysisWorkflow
from workflows.bond_comparison_workflow import BondComparisonWorkflow
from workflows.information_retrieval_workflow import InformationRetrievalWorkflow
from workflows.stock_analysis_workflow import StockAnalysisWorkflow

# 导入记忆管理系统
from memory import memory_manager
from logging_config import get_logger


class WorkflowState(TypedDict):
    """主工作流状态"""
    # 输入
    user_query: str
    thread_id: Optional[str]
    recent_history: str
    
    # 意图识别结果
    intent_type: str
    confidence: float
    bond_codes: List[str]
    stock_codes: List[str]  # 新增：股票代码
    analysis_dimensions: List[str]  # ✅ 新增：分析维度
    
    # 路由信息
    workflow_type: str  # conversation, analysis, comparison, query
    
    # 执行结果
    workflow_result: Any  # 根据工作流类型返回不同格式的结果
    conversation_result: str
    
    # ✅ 子图完整状态（用于前端格式化）
    subgraph_state: Dict[str, Any]
    
    # 记忆上下文
    context: str
    
    # 元数据
    errors: Annotated[List[str], operator.add]
    completed_steps: Annotated[List[str], operator.add]
    timestamp: str


class WorkflowManager:
    """工作流管理器 - LangGraph主图架构"""
    
    def __init__(self):
        """初始化工作流管理器"""
        self.logger = get_logger("workflows.manager")
        
        # 初始化智能体
        self.intent_agent = GeneralIntentRecognitionAgent()
        self.conversational_agent = GeneralConversationalAgent()
        
        # 初始化子工作流
        self.bond_analysis_workflow = BondAnalysisWorkflow()
        self.bond_comparison_workflow = BondComparisonWorkflow()
        self.info_retrieval_workflow = InformationRetrievalWorkflow()
        self.stock_analysis_workflow = StockAnalysisWorkflow()
        
        # 创建检查点保存器
        self.checkpointer = MemorySaver()
        
        # 构建主工作流图
        self.app = self._build_main_workflow()
        
        self.logger.info("工作流管理器初始化完成（LangGraph主图架构）")
    
    def _build_main_workflow(self):
        """构建主工作流状态图"""
        # 创建主状态图
        workflow = StateGraph(WorkflowState)
        
        # 添加节点
        workflow.add_node("context_loading", self._context_loading_node)
        workflow.add_node("intent_recognition", self._intent_recognition_node)
        workflow.add_node("conversation", self._conversation_node)
        workflow.add_node("bond_analysis", self._bond_analysis_node)
        workflow.add_node("bond_comparison", self._bond_comparison_node)
        workflow.add_node("stock_analysis", self._stock_analysis_node)
        workflow.add_node("information_retrieval", self._information_retrieval_node)
        workflow.add_node("memory_saving", self._memory_saving_node)
        
        # 设置入口点
        workflow.set_entry_point("context_loading")
        
        # 上下文加载后进行意图识别
        workflow.add_edge("context_loading", "intent_recognition")
        
        # 意图识别后根据意图类型路由
        workflow.add_conditional_edges(
            "intent_recognition",
            self._route_by_intent,
            {
                "conversation": "conversation",
                "analysis": "bond_analysis",
                "comparison": "bond_comparison",
                "stock_analysis": "stock_analysis",
                "query": "information_retrieval"
            }
        )
        
        # 所有业务节点完成后保存记忆
        workflow.add_edge("conversation", "memory_saving")
        workflow.add_edge("bond_analysis", "memory_saving")
        workflow.add_edge("bond_comparison", "memory_saving")
        workflow.add_edge("stock_analysis", "memory_saving")
        workflow.add_edge("information_retrieval", "memory_saving")
        
        # 记忆保存后结束
        workflow.add_edge("memory_saving", END)
        
        # 编译工作流
        return workflow.compile(checkpointer=self.checkpointer)
    
    async def _context_loading_node(self, state: WorkflowState) -> Dict[str, Any]:
        """上下文加载节点"""
        try:
            writer = get_stream_writer()
            
            writer({
                "type": "step_start",
                "node": "context_loading",
                "step_name": "📚 加载上下文",
                "timestamp": datetime.now().isoformat()
            })
            
            thread_id = state.get("thread_id")
            user_query = state.get("user_query", "")
            context = ""
            
            # 从记忆系统获取上下文
            if thread_id:
                context = memory_manager.get_context_for_query(thread_id, user_query)
                if context:
                    self.logger.info(f"获取到上下文信息: {len(context)} 字符")
                    writer({
                        "type": "reasoning",
                        "node": "context_loading",
                        "content": f"📋 加载对话上下文（{len(context)} 字符）",
                        "timestamp": datetime.now().isoformat()
                    })
                
                # 记录用户消息到对话记忆
                memory_manager.conversation.add_message(thread_id, "user", user_query)
            
            writer({
                "type": "step_end",
                "node": "context_loading",
                "step_name": "📚 加载上下文",
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "context": context,
                "completed_steps": ["context_loading"]
            }
        except Exception as e:
            self.logger.error(f"上下文加载失败: {e}")
            return {
                "context": "",
                "errors": [f"上下文加载失败: {str(e)}"],
                "completed_steps": ["context_loading"]
            }
    
    async def _intent_recognition_node(self, state: WorkflowState) -> Dict[str, Any]:
        """意图识别节点"""
        try:
            writer = get_stream_writer()
            
            writer({
                "type": "step_start",
                "node": "intent_recognition",
                "step_name": "🎯 意图识别",
                "timestamp": datetime.now().isoformat()
            })
            
            user_query = state.get("user_query", "")
            context = state.get("context", "")
            
            # 执行意图识别
            result = await self.intent_agent.analyze(user_query, context)
            intent_data = result.get("data", {})
            
            intent_type = intent_data.get("intent_type", "conversation")
            confidence = intent_data.get("confidence", 0.0)
            bond_codes = intent_data.get("bond_codes", [])
            stock_codes = intent_data.get("stock_codes", [])  # 新增：提取股票代码
            analysis_dimensions = intent_data.get("analysis_dimensions", [])  # ✅ 提取分析维度
            
            # 如果意图识别没有提取到债券代码，尝试从查询中提取
            if not bond_codes and intent_type in ["analysis", "comparison"]:
                bond_codes = re.findall(r'\b\d{6}\b', user_query)
                self.logger.info(f"从用户查询中提取债券代码: {bond_codes}")

            # 如果意图识别没有提取到股票代码，尝试从查询中提取
            if not stock_codes and intent_type in ["stock_analysis"]:
                stock_codes = re.findall(r'\b\d{6}\b', user_query)
                self.logger.info(f"从用户查询中提取股票代码: {stock_codes}")
            
            writer({
                "type": "reasoning",
                "node": "intent_recognition",
                "content": f"🔍 识别意图: **{intent_type}** (置信度: {confidence:.2%})",
                "timestamp": datetime.now().isoformat()
            })
            
            if bond_codes:
                writer({
                    "type": "reasoning",
                    "node": "intent_recognition",
                    "content": f"📋 提取债券代码: {', '.join(bond_codes)}",
                    "timestamp": datetime.now().isoformat()
                })
            
            writer({
                "type": "step_end",
                "node": "intent_recognition",
                "step_name": "🎯 意图识别",
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "intent_type": intent_type,
                "confidence": confidence,
                "bond_codes": bond_codes,
                "stock_codes": stock_codes,  # 新增：保存股票代码
                "analysis_dimensions": analysis_dimensions,  # ✅ 保存分析维度
                "workflow_type": intent_type,
                "completed_steps": ["intent_recognition"]
            }
        except Exception as e:
            self.logger.error(f"意图识别失败: {e}")
            return {
                "intent_type": "conversation",
                "confidence": 0.0,
                "bond_codes": [],
                "stock_codes": [],  # 新增：默认空股票代码
                "analysis_dimensions": [],  # ✅ 默认空维度
                "workflow_type": "conversation",
                "errors": [f"意图识别失败: {str(e)}"],
                "completed_steps": ["intent_recognition"]
            }
    
    def _route_by_intent(self, state: WorkflowState) -> str:
        """根据意图类型路由"""
        intent_type = state.get("intent_type", "conversation")
        
        # 对于follow_up意图，也路由到conversation
        if intent_type == "follow_up":
            return "conversation"
        
        # 映射意图类型到工作流
        route_map = {
            "conversation": "conversation",
            "analysis": "analysis",
            "comparison": "comparison",
            "stock_analysis": "stock_analysis",
            "query": "query"
        }
        
        return route_map.get(intent_type, "conversation")
    
    async def _conversation_node(self, state: WorkflowState) -> Dict[str, Any]:
        """对话处理节点"""
        try:
            writer = get_stream_writer()
            
            writer({
                "type": "step_start",
                "node": "conversation",
                "step_name": "💬 对话处理",
                "timestamp": datetime.now().isoformat()
            })
            
            user_query = state.get("user_query", "")
            thread_id = state.get("thread_id")
            
            results = []
            async for event in self.conversational_agent.stream_analyze(user_query, thread_id=thread_id):
                event_type = event.get("type")
                
                if event_type == "reasoning":
                    # 转发推理事件
                    writer({
                        "type": "reasoning",
                        "node": "conversation",
                        "content": event.get("content", ""),
                        "timestamp": datetime.now().isoformat()
                    })
                elif event_type == "final_result":
                    results.append(event.get("content", ""))
            
            conversation_result = "\n".join(results) if results else ""
            
            writer({
                "type": "step_end",
                "node": "conversation",
                "step_name": "💬 对话处理",
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "conversation_result": conversation_result,
                "workflow_result": conversation_result,  # ✅ 改为字符串
                "workflow_type": "conversation",
                "completed_steps": ["conversation"]
            }
        except Exception as e:
            self.logger.error(f"对话处理失败: {e}")
            return {
                "conversation_result": f"对话处理失败: {str(e)}",
                "errors": [f"对话处理失败: {str(e)}"],
                "completed_steps": ["conversation"]
            }
    
    async def _bond_analysis_node(self, state: WorkflowState) -> Dict[str, Any]:
        """债券分析节点（子图）- 只保存子图状态，不做格式化"""
        try:
            writer = get_stream_writer()
            
            bond_codes = state.get("bond_codes", [])
            
            if not bond_codes:
                writer({
                    "type": "error",
                    "node": "bond_analysis",
                    "error": "未提供债券代码",
                    "timestamp": datetime.now().isoformat()
                })
                return {
                    "errors": ["未提供债券代码"],
                    "completed_steps": ["bond_analysis"]
                }
            
            bond_code = bond_codes[0]
            user_query = state.get("user_query", "")
            thread_id = state.get("thread_id")
            analysis_dimensions = state.get("analysis_dimensions", [])  # ✅ 获取分析维度
            
            self.logger.info(f"开始债券分析子图: {bond_code}, 维度: {analysis_dimensions}")
            
            # ✅ 调用子图的流式执行，所有子图事件会自动传播
            subgraph_state = {}
            
            async for event in self.bond_analysis_workflow.stream_execute(
                bond_code=bond_code,
                user_query=user_query,
                thread_id=thread_id,
                analysis_dimensions=analysis_dimensions  # ✅ 传递分析维度
            ):
                # 直接转发子图事件（LangGraph会自动传播）
                writer(event)
                
                # ✅ 保存子图的完整状态（不做任何格式化）
                if event.get("type") == "final_result":
                    subgraph_state = event.get("state", {})
                    self.logger.info(f"保存债券分析子图状态，bond_code: {subgraph_state.get('bond_code')}")
            
            # ✅ 只返回子图状态和工作流类型，格式化由前端负责
            return {
                "subgraph_state": subgraph_state,
                "workflow_type": "analysis",
                "completed_steps": ["bond_analysis"]
            }
            
        except Exception as e:
            self.logger.error(f"债券分析失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "bond_analysis",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "errors": [f"债券分析失败: {str(e)}"],
                "completed_steps": ["bond_analysis"]
            }
    
    async def _stock_analysis_node(self, state: WorkflowState) -> Dict[str, Any]:
        """股票分析节点（子图）- 只保存子图状态，不做格式化"""
        try:
            writer = get_stream_writer()

            stock_codes = state.get("stock_codes", [])

            if not stock_codes:
                writer({
                    "type": "error",
                    "node": "stock_analysis",
                    "error": "未提供股票代码",
                    "timestamp": datetime.now().isoformat()
                })
                return {
                    "errors": ["未提供股票代码"],
                    "completed_steps": ["stock_analysis"]
                }

            stock_code = stock_codes[0]
            user_query = state.get("user_query", "")
            thread_id = state.get("thread_id")

            self.logger.info(f"开始股票分析子图: {stock_code}")

            # ✅ 调用子图的流式执行，所有子图事件会自动传播
            subgraph_state = {}

            async for event in self.stock_analysis_workflow.stream_execute(
                stock_code=stock_code,
                user_query=user_query,
                thread_id=thread_id,
                analysis_dimensions=state.get("analysis_dimensions", [])  # ✅ 传递分析维度
            ):
                # 直接转发子图事件（LangGraph会自动传播）
                writer(event)

                # ✅ 保存子图的完整状态（不做任何格式化）
                if event.get("type") == "final_result":
                    subgraph_state = event.get("state", {})
                    self.logger.info(f"保存股票分析子图状态，stock_code: {subgraph_state.get('stock_code')}")
                    
                    # ✅ 关键修复：详细记录final_result事件的状态
                    self.logger.info(f"🔍 final_result事件状态检查: event.state类型={type(subgraph_state)}")
                    if subgraph_state is None:
                        self.logger.error("❌ final_result事件中的state为None，这将导致后续NoneType错误")
                        subgraph_state = {}  # 防止后续NoneType错误
                    else:
                        self.logger.info(f"✅ final_result事件状态有效，包含键: {list(subgraph_state.keys())}")
                    
                    # ✅ 修复：详细记录图表数据状态
                    if "investment_charts" in subgraph_state:
                        charts_data = subgraph_state["investment_charts"]
                        if isinstance(charts_data, dict) and charts_data:
                            self.logger.info(f"检测到图表数据，包含字段: {list(charts_data.keys())}")
                            if "five_dimension_radar" in charts_data:
                                radar_data = charts_data["five_dimension_radar"]
                                if isinstance(radar_data, dict):
                                    dimensions = radar_data.get("dimensions", [])
                                    scores = radar_data.get("scores", [])
                                    self.logger.info(f"五维雷达图数据: 维度={dimensions}, 评分={scores}")
                        else:
                            self.logger.warning("investment_charts字段存在但为空或无效")
                    else:
                        self.logger.warning("子图状态中未找到investment_charts字段")

            # ✅ 只返回子图状态和工作流类型，格式化由前端负责
            return {
                "subgraph_state": subgraph_state,
                "workflow_type": "stock_analysis",
                "completed_steps": ["stock_analysis"]
            }

        except Exception as e:
            self.logger.error(f"股票分析失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "stock_analysis",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "errors": [f"股票分析失败: {str(e)}"],
                "completed_steps": ["stock_analysis"]
            }

    async def _bond_comparison_node(self, state: WorkflowState) -> Dict[str, Any]:
        """债券对比节点（子图）- 只保存子图状态，不做格式化"""
        try:
            writer = get_stream_writer()
            
            bond_codes = state.get("bond_codes", [])
            
            if len(bond_codes) < 2:
                writer({
                    "type": "error",
                    "node": "bond_comparison",
                    "error": "请提供至少2个债券代码进行对比",
                    "timestamp": datetime.now().isoformat()
                })
                return {
                    "errors": ["请提供至少2个债券代码进行对比"],
                    "completed_steps": ["bond_comparison"]
                }
            
            user_query = state.get("user_query", "")
            thread_id = state.get("thread_id")
            
            self.logger.info(f"开始债券对比子图: {bond_codes}")
            
            # ✅ 调用子图的流式执行，所有子图事件会自动传播
            subgraph_state = {}
            
            async for event in self.bond_comparison_workflow.stream_execute(
                bond_codes=bond_codes,
                user_query=user_query,
                thread_id=thread_id
            ):
                # 转发子图事件
                writer(event)
                
                # ✅ 统一捕获final_result事件（与债券分析保持一致）
                if event.get("type") == "final_result":
                    subgraph_state = event.get("state", {})
                    self.logger.info(f"保存债券对比子图状态，bond_codes: {subgraph_state.get('bond_codes')}")
            
            # ✅ 只返回子图状态和工作流类型，格式化由前端负责
            return {
                "subgraph_state": subgraph_state,
                "workflow_type": "comparison",
                "completed_steps": ["bond_comparison"]
            }
        except Exception as e:
            self.logger.error(f"债券对比失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "bond_comparison",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "errors": [f"债券对比失败: {str(e)}"],
                "completed_steps": ["bond_comparison"]
            }
    
    async def _information_retrieval_node(self, state: WorkflowState) -> Dict[str, Any]:
        """信息检索节点（子图）- 只保存子图状态，不做格式化"""
        try:
            writer = get_stream_writer()
            
            user_query = state.get("user_query", "")
            thread_id = state.get("thread_id")
            
            self.logger.info(f"开始信息检索子图: {user_query}")
            
            # ✅ 调用子图的流式执行，所有子图事件会自动传播
            subgraph_state = {}
            
            async for event in self.info_retrieval_workflow.stream_execute(
                user_query=user_query,
                thread_id=thread_id
            ):
                # 转发子图事件
                writer(event)
                
                # ✅ 统一捕获final_result事件（与债券分析保持一致）
                if event.get("type") == "final_result":
                    subgraph_state = event.get("state", {})
                    self.logger.info(f"保存信息检索子图状态，query: {user_query}")
            
            # ✅ 只返回子图状态和工作流类型，格式化由前端负责
            return {
                "subgraph_state": subgraph_state,
                "workflow_type": "query",
                "completed_steps": ["information_retrieval"]
            }
        except Exception as e:
            self.logger.error(f"信息检索失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "information_retrieval",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "errors": [f"信息检索失败: {str(e)}"],
                "completed_steps": ["information_retrieval"]
            }
    
    async def _memory_saving_node(self, state: WorkflowState) -> Dict[str, Any]:
        """记忆保存节点"""
        try:
            writer = get_stream_writer()
            
            writer({
                "type": "step_start",
                "node": "memory_saving",
                "step_name": "💾 保存记忆",
                "timestamp": datetime.now().isoformat()
            })
            
            thread_id = state.get("thread_id")
            workflow_type = state.get("workflow_type", "")
            
            if not thread_id:
                writer({
                    "type": "step_end",
                    "node": "memory_saving",
                    "step_name": "💾 保存记忆",
                    "timestamp": datetime.now().isoformat()
                })
                return {"completed_steps": ["memory_saving"]}
            
            # 保存助手回复到对话记忆
            assistant_response = self._extract_assistant_response(state)
            if assistant_response:
                memory_manager.conversation.add_message(thread_id, "assistant", assistant_response)
                writer({
                    "type": "reasoning",
                    "node": "memory_saving",
                    "content": "📝 保存对话记录",
                    "timestamp": datetime.now().isoformat()
                })
            
            # 保存分析结果到记忆
            if workflow_type in ["analysis", "comparison", "stock_analysis"]:
                self._save_analysis_to_memory(state, thread_id)
                writer({
                    "type": "reasoning",
                    "node": "memory_saving",
                    "content": f"💾 保存{workflow_type}结果",
                    "timestamp": datetime.now().isoformat()
                })
            
            writer({
                "type": "step_end",
                "node": "memory_saving",
                "step_name": "💾 保存记忆",
                "timestamp": datetime.now().isoformat()
            })
            
            return {"completed_steps": ["memory_saving"]}
        except Exception as e:
            self.logger.error(f"记忆保存失败: {e}")
            return {
                "errors": [f"记忆保存失败: {str(e)}"],
                "completed_steps": ["memory_saving"]
            }
    
    def _extract_assistant_response(self, state: WorkflowState) -> str:
        """从状态中提取助手回复内容"""
        workflow_type = state.get("workflow_type", "")
        conversation_result = state.get("conversation_result", "")
        workflow_result = state.get("workflow_result")  # ✅ 支持任意类型
        subgraph_state = state.get("subgraph_state", {})
        
        # ✅ 关键修复：详细记录状态数据完整性
        self.logger.info(f"🔍 _extract_assistant_response详细检查:")
        self.logger.info(f"  - workflow_type: {workflow_type}")
        self.logger.info(f"  - conversation_result: {bool(conversation_result)}")
        self.logger.info(f"  - workflow_result: {workflow_result is not None}")
        self.logger.info(f"  - subgraph_state: {subgraph_state is not None}")
        
        if subgraph_state is not None:
            self.logger.info(f"  - subgraph_state类型: {type(subgraph_state)}")
            if isinstance(subgraph_state, dict):
                self.logger.info(f"  - subgraph_state包含键: {list(subgraph_state.keys())}")
            else:
                self.logger.warning(f"  - subgraph_state不是字典类型: {type(subgraph_state)}")
        else:
            self.logger.warning("  - subgraph_state为None，这可能导致后续NoneType错误")
        
        # 对话结果
        if conversation_result:
            self.logger.info("✅ 使用conversation_result作为助手回复")
            return conversation_result
        
        # ✅ 债券分析结果 - 从subgraph_state提取
        if workflow_type == "analysis" and subgraph_state:
            # 提取最终总结
            final_summary = subgraph_state.get("final_summary", "")
            if final_summary:
                bond_code = subgraph_state.get("bond_code", "")
                self.logger.info(f"✅ 使用债券分析结果作为助手回复: {bond_code}")
                return f"已完成债券{bond_code}的分析。{final_summary}"
        
        # ✅ 股票分析结果 - 从subgraph_state提取
        if workflow_type == "stock_analysis" and subgraph_state:
            # ✅ 关键修复：确保subgraph_state不为None
            if subgraph_state is None:
                self.logger.error("❌ subgraph_state为None，无法提取股票分析结果")
                return "股票分析完成，但结果数据异常"
            
            # 提取投资总结
            investment_summary = subgraph_state.get("investment_summary", "")
            if investment_summary:
                stock_code = subgraph_state.get("stock_code", "")
                self.logger.info(f"✅ 使用股票分析结果作为助手回复: {stock_code}")
                return f"已完成股票{stock_code}的分析。{investment_summary}"
        
        # ✅ 对比和检索结果 - 从subgraph_state提取
        if workflow_type in ["comparison", "query"] and subgraph_state:
            # 尝试获取格式化结果
            formatted_result = subgraph_state.get("formatted_result", "")
            if formatted_result:
                self.logger.info("✅ 使用对比/检索结果作为助手回复")
                return formatted_result
        
        # 旧版workflow_result格式
        if workflow_result:
            self.logger.info("✅ 使用workflow_result作为助手回复")
            return workflow_result  # ✅ 直接返回字符串
        
        self.logger.warning("⚠️ 无法提取助手回复，返回空字符串")
        return ""
    
    def _save_analysis_to_memory(self, state: WorkflowState, thread_id: str):
        """保存分析结果到记忆系统"""
        try:
            workflow_type = state.get("workflow_type")
            bond_codes = state.get("bond_codes", [])
            stock_codes = state.get("stock_codes", [])

            codes = bond_codes or stock_codes
            if not codes:
                return

            analysis_data = {
                "workflow": workflow_type,
                "user_query": state.get("user_query", ""),
                "result": state.get("workflow_result"),  # ✅ 支持任意类型
                "timestamp": state.get("timestamp", datetime.now().isoformat())
            }

            # 对于分析工作流
            if workflow_type == "analysis":
                bond_code = codes[0]
                memory_manager.save_analysis_result(
                    bond_code=bond_code,
                    analysis_type="comprehensive",
                    result=analysis_data,
                    thread_id=thread_id
                )
                self.logger.info(f"已保存债券 {bond_code} 的分析结果到记忆系统")

            # 对于股票分析工作流
            elif workflow_type == "stock_analysis":
                stock_code = codes[0]
                memory_manager.save_analysis_result(
                    bond_code=stock_code,  # 复用bond_code字段，实际是股票代码
                    analysis_type="stock_comprehensive",
                    result=analysis_data,
                    thread_id=thread_id
                )
                self.logger.info(f"已保存股票 {stock_code} 的分析结果到记忆系统")

            # 对于对比工作流
            elif workflow_type == "comparison":
                for code in codes[:5]:
                    memory_manager.save_analysis_result(
                        bond_code=code,
                        analysis_type="comparison",
                        result=analysis_data,
                        thread_id=thread_id
                    )
                self.logger.info(f"已保存 {len(codes)} 个代码的对比结果到记忆系统")

        except Exception as e:
            self.logger.error(f"保存分析结果到记忆系统失败: {e}")
    
    async def stream_execute(
        self, 
        user_query: str, 
        thread_id: Optional[str] = None,
        recent_history: str = ""
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        流式执行主工作流（LangGraph主图+子图架构）
        
        ✅ 完整的事件流传播：
        - 主图节点事件（上下文加载、意图识别、记忆保存）
        - 子图所有事件（债券分析、对比、检索的所有节点和推理过程）
        - 实现真正的端到端可观测性
        
        Args:
            user_query: 用户查询
            thread_id: 线程ID
            recent_history: 最近的对话历史
            
        Yields:
            主图和子图的所有执行事件
        """
        if not thread_id:
            thread_id = f"workflow_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        # 初始化状态
        initial_state: WorkflowState = {
            "user_query": user_query,
            "thread_id": thread_id,
            "recent_history": recent_history,
            "intent_type": "",
            "confidence": 0.0,
            "bond_codes": [],
            "stock_codes": [],  # 新增：股票代码字段
            "analysis_dimensions": [],  # ✅ 添加分析维度字段
            "workflow_type": "",
            "workflow_result": None,  # ✅ 初始化为None，支持任意类型
            "conversation_result": "",
            "subgraph_state": {},  # ✅ 添加子图状态字段
            "context": "",
            "errors": [],
            "completed_steps": [],
            "timestamp": datetime.now().isoformat()
        }
        
        config = RunnableConfig(configurable={"thread_id": thread_id})
        
        try:
            # ✅ 使用stream_mode="custom"接收主图和子图的所有自定义事件
            async for event in self.app.astream(
                initial_state, 
                config, 
                stream_mode="custom"
            ):
                # 直接转发所有事件（主图和子图）
                yield event
            
            # 获取最终状态
            final_state = await self.app.aget_state(config)
            
            # 发送最终结果事件
            yield {
                "type": "workflow_complete",
                "state": final_state.values,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error(f"流式执行失败: {e}", exc_info=True)
            yield {
                "type": "error",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
    
    def get_available_workflows(self) -> List[str]:
        """获取可用的工作流列表"""
        return ["conversation", "analysis", "comparison", "stock_analysis", "query"]
    
    def get_workflow_info(self, workflow_name: str) -> Dict[str, Any]:
        """获取工作流信息"""
        workflow_info = {
            "conversation": {
                "name": "对话工作流",
                "description": "处理一般性对话和问答",
                "intent_types": ["conversation", "follow_up"]
            },
            "analysis": {
                "name": "债券分析工作流",
                "description": "完整的债券分析，包括信用、条款、宏观、估值等维度",
                "intent_types": ["analysis"]
            },
            "comparison": {
                "name": "债券对比工作流",
                "description": "对比分析多个债券（2-5个）",
                "intent_types": ["comparison"]
            },
            "stock_analysis": {
                "name": "股票分析工作流",
                "description": "完整的股票分析，包括技术、消息、估值、资金四大维度",
                "intent_types": ["stock_analysis"]
            },
            "query": {
                "name": "信息检索工作流",
                "description": "检索债券、公告、新闻等信息",
                "intent_types": ["query"]
            }
        }
        
        return workflow_info.get(workflow_name, {
            "name": "未知工作流",
            "description": "工作流不存在"
        })


# 全局实例
workflow_manager = WorkflowManager()


# 便捷函数
async def process_query(user_query: str, thread_id: Optional[str] = None) -> AsyncIterator[Dict[str, Any]]:
    """
    处理用户查询的便捷函数（流式）
    
    Args:
        user_query: 用户查询
        thread_id: 线程ID
        
    Yields:
        处理过程的所有事件
    """
    async for event in workflow_manager.stream_execute(user_query, thread_id):
        yield event